#include "util.hpp"
#include "perm_parser.hpp"
#include "snp_cand.hpp"
#include "snp_call.hpp"
#include "call_rates.hpp"
#include "consensus.hpp"
#include "easyrun.hpp"
#include "view.hpp"
#include "priors.hpp"

#include <utility>

#include <boost/program_options.hpp>
#include <boost/lexical_cast.hpp>
namespace po = boost::program_options;

/** Set all default options for comb. programs
 */
void options_set_defaults(prog_options &options)
{
    options.consec_snps          = 3;
    options.diploid              = false;
    options.has_cand             = false;
    options.ignore_missing       = false;
    options.max_block_size       = 5;
    // TODO: default 36? should it be an option, or determined from file?
    //options.nbins                = 36;
    options.normalize_freq       = 10;
    options.old_bin_cnt          = 0;
    //options.old_bin_cnt          = 36; // color
    //options.old_bin_cnt          = 96; // sequence
    options.new_bin_cnt          = 10;
    options.priors_str           = "M";
    options.read_count           = 10000000;
    options.max_map_loci         = 25;
    options.unique               = false;
    options.expr_all             = false;
    options.expr_unique          = false;
    options.expr_random          = false;
    options.only_show_changed    = false;
}

/** Load priors for given problem. */
void load_priors(prog_options &options, po::variables_map vm)
{
    if( !options.diploid )
    {
        if(vm.count("priors-file"))
        {
            std::ifstream ifs(options.priors_file.c_str(), std::ios_base::in);
            if(!ifs)
                error("Could not open file: " + options.priors_file);
            load_haploid_priors_from_file(ifs, options.haploid_priors);
            ifs.close();
        }
        else
        {
            if(options.priors_str == "N")
                load_haploid_priors_builtin(haploid_priors_none, options.haploid_priors);
            else if(options.priors_str == "H")
                load_haploid_priors_builtin(haploid_priors_high, options.haploid_priors);
            else if(options.priors_str == "M")
                load_haploid_priors_builtin(haploid_priors_moderate, options.haploid_priors);
            else
                error("Unable to recognize priors option: " + options.priors_str + " expected N (none), M (medium), or H (high)");
        }
    }
    else // diploid case.
    {
        if(vm.count("priors-file"))
        {
            std::ifstream ifs(options.priors_file.c_str(), std::ios_base::in);
            if(!ifs)
                error("Could not open file: " + options.priors_file);
            load_diploid_priors_from_file(ifs, options.diploid_priors);
            ifs.close();
        }
        else // priors
        {
            if(options.priors_str == "N")
                load_diploid_priors_builtin(diploid_priors_none, options.diploid_priors);
            else if(options.priors_str == "M")
                load_diploid_priors_builtin(diploid_priors_moderate, options.diploid_priors);
            else if(options.priors_str == "H")
                load_diploid_priors_builtin(diploid_priors_high, options.diploid_priors);
            else
                error("Unable to recognize priors option: " + options.priors_str + " expected N (none), M (medium), or H (high)");
        }
    }
}

prog_options parse_expression_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    std::string usage = "Usage: comb expression [options] reference.fasta mapping.out";

    po::options_description visible("Command line options");

    visible.add_options()
        ("help", "Produce help message")
        ("all", po::bool_switch(&options.expr_all), "Map reads to all locations")
        ("random", po::bool_switch(&options.expr_random), "Map reads to a location at random")
        ("unique", po::bool_switch(&options.expr_unique), "calculate expression using only uniquely mapping reads [DEFAULT]")
        ("call-file", po::value<std::string>(&options.call_file), "Call rates file necessary for weighting (--random or --weighted)")
    ;

    po::options_description hidden;

    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "Genome sequence.")
        ("mapping-file", po::value<std::string>(&options.map_file), "Mapping file")
    ;
        
    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("mapping-file", 1);
    
    po::variables_map vm;
        
    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    if( vm["all"].as<bool>() + vm["unique"].as<bool>() + vm["random"].as<bool>() == 0)
        options.expr_unique = true;
    //else if( vm.count("all") + vm.count("unique") + vm.count("random") > 1)
    else if( vm["all"].as<bool>() + vm["unique"].as<bool>() + vm["random"].as<bool>() > 1)
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Please specify only one of unqiue/all/random");
    }

    /* Make sure all required files are given. */
    if(!vm.count("mapping-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional arguments given.");
    }
    
    return options;
}


prog_options parse_consensus_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    std::string usage = "Usage: comb consensus [options] reference.fasta comb_snp.out";

    po::options_description visible("Command line options");

    visible.add_options()
        ("help", "Produce help message")
        ("diploid", po::bool_switch(&options.diploid), "Call using diploid models. [Default false]")
        ("priors-file", po::value<std::string>(&options.priors_file), "Load priors from file.")
        ("priors", po::value<std::string>(&options.priors_str), "Choose a set of priors [DEFAULT M (moderate)]")
    ;

    po::options_description hidden;

    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "Genome sequences file")
        ("snps-file", po::value<std::string>(&options.snp_file), "Comb SNPS file")
    ;
        
    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("snps-file", 1);
    
    po::variables_map vm;
        
    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    /* Make sure all required files are given. */
    if(!vm.count("snps-file")) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional arguments given.");
    }

    load_priors(options, vm);

    return options;
}


prog_options parse_callrates_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);
    options.ignore_missing = true;

    std::string usage = "Usage: comb callrates [options] reference.fasta mapping.out";

    po::options_description visible("Command line options");
    
    visible.add_options()
        ("help,h", "show this help message")
        ("count,c", po::value<unsigned int>(&options.read_count), "Number of reads to use [DEFAULT 10,000,000]")
        ("unique,u", po::bool_switch(&options.unique), "Use unique reads.")
        ("snp-cand", po::value<std::string>(&options.cand_file), "SNP Cand File")
    ;

    po::options_description hidden;

    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "Genome sequence.")
        ("map-file", po::value<std::string>(&options.map_file), "Mapping File")
    ;

    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("map-file", 1);
        
    po::variables_map vm;

    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    /* Make sure all required files are given. */
    if(!vm.count("map-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional arguments.");
    }
    
    if(vm.count("snp-cand")) 
    {
        options.has_cand = true;
    }
    return options;
}

prog_options parse_mapcount_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    po::options_description visible("Command line options");
    std::string usage = "Usage: comb mapcount [options] reference.fasta mapping.out";

    visible.add_options()
        ("help", "Produce help message")
        ("unique", po::bool_switch(&options.unique), "Use only uniquely mapping reads")
        ("ignore-missing", po::bool_switch(&options.ignore_missing), "Allow mapping locations not in reference.")
        ;

    po::options_description hidden;

    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "FASTA formatted genome")
        ("map-file", po::value<std::string>(&options.map_file), "Mapping File")
    ;

    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("map-file", 1);
        
    po::variables_map vm;

    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }


    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    /* Make sure all required files are given. */
    if(!vm.count("map-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional arguments given.");
    }
    return options;
}

prog_options parse_mapcand_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    po::options_description visible("Command line options");
    std::string usage = "Usage: comb mapcand [options] reference.fasta cand_mapcount.out";

    std::vector<std::string> cover_accum;
    visible.add_options()
        ("help", "Produce help message")
        ("cover", po::value<std::vector<std::string> >(&cover_accum), "cover,threshold pairs")
        ("max-block", po::value<unsigned int>(&options.max_block_size), "The largest block of consecutive snps to allow. [DEFAULT 5]")
    ;

    po::options_description hidden;

    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "genome reference.")
        ("mapcount-file", po::value<std::string>(&options.map_count_file), "File of mapping statistics to analyze.")
    ;

    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("mapcount-file", 1);
    
    po::variables_map vm;

    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    /* Make sure all required files are given. */
    if(!vm.count("mapcount-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional args. given");
    }
    
    /* Parse new cover format. */
    if(cover_accum.size() == 0)
    {
        options.cover_threshold.push_back(std::make_pair<int, float>(1, 0.5));
        options.cover_threshold.push_back(std::make_pair<int, float>(4, 0.2));
    }

    for(unsigned int i = 0; i < cover_accum.size(); i++)
    {
        std::vector<std::string> toks;
        strtok(cover_accum[i], ",", toks);
        if(toks.size() != 2)
            error("Failed to parse cover argument as a pair of comma seperated values: " + cover_accum[i]);
        int   c = boost::lexical_cast<int>(toks[0]);
        float t = boost::lexical_cast<float>(toks[1]);

        options.cover_threshold.push_back(std::make_pair<int, float>(c,t));
    }

    return options;
}


prog_options parse_snpcand_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    po::options_description visible("Command line options"); 
    std::string usage = "Usage: comb snpcand [options] reference.fasta mapping.out";
    
    std::vector<std::string> cover_accum;
    visible.add_options()
        ("help", "Produce help message")
        ("cover", po::value<std::vector<std::string> >(&cover_accum), "cover,threshold pairs")
        ("unique", po::bool_switch(&options.unique), "Look for unique reads")
        ("ignore-missing", po::bool_switch(&options.ignore_missing), "Allow mapping locations not in reference.")
        ("max-loci", po::value<unsigned int>(&options.max_map_loci), "Maximum mapping locations allowed.  Reads with more will be discarded. [DEFAULT 25]")
        ("max-block", po::value<unsigned int>(&options.max_block_size), "The largest block of consecutive snps to allow. [DEFAULT 5]")
    ;

    po::options_description hidden;

    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "FASTA formatted reference")
        ("map-file", po::value<std::string>(&options.map_file), "Mapping File")
    ;

    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("map-file", 1);
        
    po::variables_map vm;

    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    /* Make sure all required files are given. */
    if(!vm.count("map-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional arguments given.");
    }
    
    /* Parse new cover format. */
    if(cover_accum.size() == 0)
    {
        options.cover_threshold.push_back(std::make_pair<int, float>(1, 0.5));
        options.cover_threshold.push_back(std::make_pair<int, float>(4, 0.2));
    }

    for(unsigned int i = 0; i < cover_accum.size(); i++)
    {
        std::vector<std::string> toks;
        strtok(cover_accum[i], ",", toks);
        if(toks.size() != 2)
            error("Failed to parse cover argument as a pair of comma seperated values: " + cover_accum[i]);
        int   c = boost::lexical_cast<int>(toks[0]);
        float t = boost::lexical_cast<float>(toks[1]);

        options.cover_threshold.push_back(std::make_pair<int, float>(c,t));
    }

    return options;
}

prog_options parse_snpcall_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    po::options_description visible("Command line options");
    std::string usage =  "Usage: comb snpcall [options] reference.fasta comb_snpcand.out mapping.out";

    visible.add_options()
        ("help,h", "Print this help message")
        ("unique", po::bool_switch(&options.unique), "Use only reads that map uniquely.")
        ("diploid", po::bool_switch(&options.diploid), "Call diploid")
        ("freq,f", po::value<int>(&options.normalize_freq), "Normalization frequency (Only needed if NaNs appear in SNP call output).")
        ("call-rates,r", po::value<std::string>(&options.call_file), "Use processed read quality information in SNP calls.")
        ("ignore-missing", po::bool_switch(&options.ignore_missing), "Allow ignore_missing mappings, mapping locations not in the reference.")
    ;

    po::options_description hidden;
    
    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "Fasta File")
        ("snp-cand", po::value<std::string>(&options.cand_file), "SNP Cand File")
        ("map-file", po::value<std::string>(&options.map_file), "Mapping File")
    ;

    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("snp-cand", 1);
    p.add("map-file", 1);
        
    po::variables_map vm;

    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 3 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    /* Make sure all required files are given. */
    if(!vm.count("map-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional arguments given.");
    }

    if( options.normalize_freq < 1 )
        error("normalize_freq must be > 1.");

    load_priors(options, vm);
    return options;
}
    
prog_options parse_easyrun_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    po::options_description visible("Command line options");
    std::string usage = "Usage: comb easyrun [options] reference.fasta mapping.out";

    std::vector<std::string> cover_accum;

    visible.add_options()
        ("help", "Produce help message")
        ("count,c", po::value<unsigned int>(&options.read_count), "Number of reads to use in call rates calculation [DEFAULT 10,000,000]")
        ("unique,u", po::bool_switch(&options.unique), "Use uniquely mapping reads.")
        ("ignore-missing", po::bool_switch(&options.ignore_missing), "Allow mapping locations not in reference.")
        ("cover", po::value<std::vector<std::string> >(&cover_accum), "cover,threshold pairs")
        ("max-block", po::value<unsigned int>(&options.max_block_size), "The largest block of consecutive snps to allow. [DEFAULT 5]")
        ("diploid", po::bool_switch(&options.diploid), "Turn on 'Diploid'")
        ("freq,f", po::value<int>(&options.normalize_freq), "Normalization frequency (Only needed if NaNs appear in SNP call output).")
        ("priors-file", po::value<std::string>(&options.priors_file), "Load priors from file.")
        ("priors", po::value<std::string>(&options.priors_str), "Choose a set of priors [DEFAULT M (moderate)]")
        ("consensus", po::value<std::string>(&options.cons_file), "Write out a consensus genome.")
        ("hide-unchanged", po::bool_switch(&options.only_show_changed), "Only show called SNPs that are different then the original base. [DEFAULT show all snps]")
    ;

    po::options_description hidden;

    hidden.add_options()
        ("fasta-file", po::value<std::string>(&options.fasta_file), "FASTA genome sequence file")
        ("map-file", po::value<std::string>(&options.map_file), "Mapping File")
    ;

    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("fasta-file", 1);
    p.add("map-file", 1);
        
    po::variables_map vm;

    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }

    /* Make sure all required files are given. */
    if(!vm.count("map-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Not enough positional arguments given.");
    }
   
    /* Parse new cover format. */
    if(cover_accum.size() == 0)
    {
        options.cover_threshold.push_back(std::make_pair<int, float>(1, 0.5));
        options.cover_threshold.push_back(std::make_pair<int, float>(4, 0.2));
    }

    for(unsigned int i = 0; i < cover_accum.size(); i++)
    {
        std::vector<std::string> toks;
        strtok(cover_accum[i], ",", toks);
        if(toks.size() != 2)
            error("Failed to parse cover argument as a pair of comma seperated values: " + cover_accum[i]);
        int   c = boost::lexical_cast<int>(toks[0]);
        float t = boost::lexical_cast<float>(toks[1]);

        options.cover_threshold.push_back(std::make_pair<int, float>(c,t));
    }

    if( options.normalize_freq < 1 )
        error("normalize_freq must be > 1.");

    load_priors(options, vm);
    return options;
}

prog_options parse_view_options(int argc, char **argv)
{
    prog_options options;
    options_set_defaults(options);

    po::options_description visible("Command line options");
    std::string usage = "Usage: comb view [options] snpcand-blob-file OR snpcall-blob-file OR mapcount-blob-file OR expression-blob-file OR ...";
    
    visible.add_options()
        ("help,h", "Print this help message")
        ("priors-file", po::value<std::string>(&options.priors_file), "Load priors from file.")
        ("priors", po::value<std::string>(&options.priors_str), "Choose a set of priors [DEFAULT M (moderate)]")
        ("hide-unchanged", po::bool_switch(&options.only_show_changed), "Only show called SNPs that are different then the original base. [DEFAULT show all snps]")
    ;

    po::options_description hidden;

    hidden.add_options()
        ("blob-file", po::value<std::string>(&options.generic_file), "File to view")
    ;

    po::options_description desc;

    desc.add(visible).add(hidden);

    po::positional_options_description p;
    p.add("blob-file", 1);
        
    po::variables_map vm;

    try 
    {
        po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
        po::notify(vm);
    }
    catch(std::exception &se) 
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("Exception caught! " + std::string(se.what()));
    }
        
    if(argc < 2 || strcmp(argv[1], "help") == 0 || vm.count("help"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        exit(0);
    }


    /* Make sure all required files are given. */
    if(!vm.count("blob-file"))
    {
        std::cerr << usage << std::endl;
        std::cerr << visible << std::endl;
        error("A blob file to view is a required argument.");
    }

    load_priors(options, vm);

    return options;
}
